package com.qingzhuge.web.security;

import cn.hutool.core.codec.Base64;
import cn.hutool.crypto.SecureUtil;
import cn.hutool.crypto.asymmetric.KeyType;
import cn.hutool.crypto.asymmetric.RSA;
import com.alibaba.fastjson.JSON;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.qingzhuge.common.annotation.SecurityParameter;
import com.qingzhuge.controller.security.ControllerSecurityBody;
import com.qingzhuge.controller.security.RsaKey;
import com.qingzhuge.dto.response.ResponseBodyDto;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.MethodParameter;
import org.springframework.http.MediaType;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.lang.Nullable;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.bind.annotation.RestControllerAdvice;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;

import javax.annotation.Resource;

/**
 * @author : zero.xiao
 * @description :
 * @date :2019-08-12 16:56
 * @modified :
 */
@Slf4j
@RestControllerAdvice(annotations = RestController.class)
public class MessageResponseBodyAdvice implements ResponseBodyAdvice<Object> {
    @Resource
    private ControllerSecurityBody controllerSecurityBody;
    /**
     * response
     */
    @Override
    public boolean supports(@Nullable MethodParameter methodParameter, @Nullable Class<? extends HttpMessageConverter<?>> converterType) {
        //获取当前处理请求的controller的方法
        if (controllerSecurityBody.isEncrypt() && null != methodParameter) {
            SecurityParameter annotation = methodParameter.getMethodAnnotation(SecurityParameter.class);
            boolean encode = null != annotation && annotation.outEncode();
            if (encode) {
                log.debug("Response needs to be encrypted : [{}]", methodParameter);
            }
            return encode;
        }
        // 不拦截/不需要处理返回值 的方法
        return false;
    }

    @Override
    public Object beforeBodyWrite(Object body,
                                  @Nullable MethodParameter methodParameter,
                                  @Nullable MediaType selectedContentType,
                                  @Nullable Class<? extends HttpMessageConverter<?>> selectedConverterType,
                                  @Nullable ServerHttpRequest request,
                                  @Nullable ServerHttpResponse response) {
        if (null != methodParameter && null != body) {
            //获取注解配置的包含和去除字段
            SecurityParameter serializedField = methodParameter.getMethodAnnotation(SecurityParameter.class);
            //入参是否需要解密
            if (null != serializedField) {
                //是否解密
                boolean encode = serializedField.outEncode();
                ResponseBodyDto.ResponseBodyDtoBuilder<Object> builder = ResponseBodyDto.builder();
                if (encode && body instanceof ResponseBodyDto) {
                    log.info("对方法method :【data】数据进行加密!");
                    String data = JSON.toJSONString(((ResponseBodyDto) body).getData());
                    log.info("开始对返回值进行加密操作!");
                    String encryptedBody = encryptedBody(data);
                    builder.message(((ResponseBodyDto) body).getMessage())
                            .code(((ResponseBodyDto) body).getCode())
                            .data(encryptedBody);
                    return builder.build();
                }
            }
        }
        return body;
    }

    private String encryptedBody(Object body) {
        try {
            ObjectMapper objectMapper = new ObjectMapper();
            String result = objectMapper.writerWithDefaultPrettyPrinter().writeValueAsString(body);
            log.info("对字符串开始加密!");
            try {
                byte[] dataBates = encrypt(result);
                String outputData = null;
                if (null != dataBates && dataBates.length > 0) {
                    outputData = Base64.encode(dataBates);
                }
                log.info("操作结束!");
                return outputData;
            } catch (Exception e) {
                e.printStackTrace();
                throw new RuntimeException("对返回数据加密出现异常：" + e.getMessage());
            }

        } catch (JsonProcessingException e) {
            e.printStackTrace();
            throw new RuntimeException("获取返回值出现异常:" + e.getMessage());
        }
    }

    private byte[] encrypt(String inputData) {
        RsaKey rsaKey = controllerSecurityBody.get();
        byte[] bytes = new byte[]{};
        if (null != rsaKey){
            RSA rsa = SecureUtil.rsa(rsaKey.getPrivateKey(), rsaKey.getPublicKey());
            bytes = rsa.encrypt(inputData, KeyType.PrivateKey);
        }
        return bytes;
    }
}
